pub mod board;
pub mod log;
pub mod parse;

use std::{
    collections::HashMap,
    fs::File,
    io::{self, Write},
    path::Path,
};

use ndarray::{Array2, ArrayView2};
// use plotters::prelude::*;
use rand::prelude::*;

use board::{BoardState, Player};

/// 棋盘大小 (19x19)
#[allow(unused)]
const BOARD_SIZE: usize = 19;

/// Q-learning 代理
#[derive(Debug)]
pub struct QLearningAgent {
    size: usize,
    epsilon: f64,                                          // 探索率
    alpha: f64,                                            // 学习率
    gamma: f64,                                            // 折扣因子
    q_table: HashMap<Array2<Option<Player>>, Array2<f64>>, // Q表: 状态 -> 动作值
}

impl QLearningAgent {
    pub fn new(size: usize, epsilon: f64, alpha: f64, gamma: f64) -> Self {
        Self {
            size,
            epsilon,
            alpha,
            gamma,
            q_table: HashMap::new(),
        }
    }

    /// 获取当前状态的Q值表
    pub fn get_q_values(&mut self, state: &ArrayView2<'_, Option<Player>>) -> Array2<f64> {
        // 如果状态不在Q表中，初始化一个新表
        self.q_table
            .entry(state.to_owned())
            .or_insert_with(|| Array2::zeros((self.size, self.size)))
            .clone()
    }

    /// 根据当前状态选择动作
    pub fn choose_action(&mut self, state: &BoardState) -> (usize, usize) {
        let valid_moves = state.valid_moves();

        if valid_moves.is_empty() {
            panic!("No valid moves available");
        }

        let mut rng = rand::rng();

        // ε-贪婪策略：以ε概率随机选择，否则选择最佳动作
        if rng.random_range(0.0..1.0) < self.epsilon {
            // 随机选择合法动作
            return valid_moves
                .iter()
                .nth(rng.random_range(0..valid_moves.len()))
                .unwrap()
                .to_owned();
        }
        // 选择Q值最大的合法动作
        let q_values = self.get_q_values(&state.board_view());
        let mut best_value = f64::NEG_INFINITY;
        let mut best_moves = Vec::new();

        for &(r, c) in valid_moves {
            let value = q_values[[r, c]];
            if value > best_value {
                best_value = value;
                best_moves.clear();
                best_moves.push((r, c));
            } else if (value - best_value).abs() < f64::EPSILON {
                best_moves.push((r, c));
            }
        }

        // 找不到最佳动作？随机返回一个
        if best_moves.is_empty() {
            return valid_moves
                .iter()
                .nth(rng.random_range(0..valid_moves.len()))
                .unwrap()
                .to_owned();
        }

        // 如果有多个相同Q值的动作，随机选择一个
        best_moves[rng.random_range(0..best_moves.len())]
    }

    /// 更新Q表
    pub fn update_q_table(
        &mut self,
        state: &ArrayView2<'_, Option<Player>>,
        action: (usize, usize),
        reward: f64,
        next_state: &ArrayView2<'_, Option<Player>>,
    ) {
        let current_q = self.get_q_values(state);
        let next_q = self.get_q_values(next_state);

        // 计算目标Q值
        let max_next_q = next_q.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));

        let target = reward + self.gamma * max_next_q;
        let td_error = target - current_q[[action.0, action.1]];

        // 更新Q值
        if let Some(q_vals) = self.q_table.get_mut(&state.to_owned()) {
            q_vals[[action.0, action.1]] += self.alpha * td_error;
        }
    }

    /// 保存Q表到文件
    #[allow(unused)]
    pub fn save_q_table(&self, file_path: &Path) -> io::Result<()> {
        todo!();
        // let mut file = File::create(file_path)?;
        // let content = bincode::encode_to_vec(&self.q_table, bincode::config::standard())
        //     .map_err(io::Error::other)?;
        // file.write_all(&content)?;
        Ok(())
    }

    /// 从文件加载Q表
    #[allow(unused)]
    pub fn load_q_table(&mut self, file_path: &Path) -> io::Result<()> {
        todo!();
        let file = File::open(file_path)?;
        // let mut reader = std::io::BufReader::new(file);
        // self.q_table = bincode::decode_from_std_read(&mut reader, bincode::config::standard())
        //     .map_err(io::Error::other)?;
        Ok(())
    }
}

/// 游戏环境
pub struct Environment<const WIN_COND: usize = 5> {
    board: BoardState<WIN_COND>,
    size: usize,
    rewards: [f64; WIN_COND],
}

pub struct StepResult {
    before: Array2<Option<Player>>,
    reward: f64,
    game_over: bool,
}

impl<const WIN_COND: usize> Environment<WIN_COND> {
    pub fn new(size: usize, rewards: [f64; WIN_COND]) -> Self {
        Self {
            board: BoardState::new(size, size),
            size,
            rewards,
        }
    }

    /// 重置环境
    pub fn reset(&mut self) {
        self.board = BoardState::new(self.size, self.size);
    }

    /// 执行动作
    pub fn step(&mut self, action: (usize, usize)) -> StepResult {
        let (row, col) = action;
        let mut reward = 0.1; // 默认合法动作奖励
        let before: Array2<_> = self.board.board_view().to_owned();

        let Ok(_) = self.board.make_move(row, col) else {
            // 非法动作惩罚
            return StepResult {
                before,
                reward: -0.1,
                game_over: false,
            };
        };

        // 检查胜利
        if self.board.check_win(Player::Black) || self.board.check_win(Player::White) {
            return StepResult {
                before,
                reward: self.rewards.last().unwrap_or(&10.0).to_owned(),
                game_over: true,
            };
        }

        // 检查平局
        if self.board.check_draw() {
            reward = 0.5; // 平局奖励
            return StepResult {
                before,
                reward,
                game_over: true,
            };
        }

        // 根据棋盘状态获取奖励
        for count in 1..=WIN_COND {
            let c_reward = self.rewards[count - 1];
            reward += c_reward * self.board.get_streak(Player::Black, count) as f64;
            reward -= c_reward * self.board.get_streak(Player::White, count) as f64;
        }

        StepResult {
            before,
            reward,
            game_over: false,
        }
    }
}

/// 训练代理
pub fn train_agent(
    agent: &mut QLearningAgent,
    env: &mut Environment,
    episodes: usize,
    save_interval: usize,
    output_dir: &Path,
) -> Vec<f64> {
    let mut total_rewards = Vec::new();

    for episode in 0..episodes {
        env.reset();
        let mut done = false;
        let mut episode_reward = 0.0;

        while !done {
            // 选择动作
            let action = agent.choose_action(&env.board);

            // 执行动作
            let StepResult {
                before,
                reward,
                game_over,
            } = env.step(action);
            done = game_over;
            episode_reward += reward;

            // 更新Q表
            agent.update_q_table(&before.view(), action, reward, &env.board.board_view());
        }

        total_rewards.push(episode_reward);

        // 定期保存Q表
        if (episode + 1) % save_interval == 0 {
            let file_name = format!("q_table_episode_{}.bin", episode + 1);
            if let Err(e) = agent.save_q_table(&output_dir.join(Path::new(&file_name))) {
                eprintln!("Failed to save Q-table: {}", e);
            } else {
                println!("Q-table saved to {}", file_name);
            }
        }

        // 打印训练进度
        if (episode + 1) % 10 == 0 {
            println!(
                "Episode {}/{}: Total Reward = {:.2}",
                episode + 1,
                episodes,
                episode_reward
            );
        }
    }

    total_rewards
}

/// 绘制奖励曲线并保存数据
pub fn save_rewards(rewards: &[f64], rewards_path: &Path) {
    // 1. 计算奖励统计信息
    let avg_reward: f64 = rewards.iter().sum::<f64>() / rewards.len() as f64;
    let min_reward = rewards.iter().fold(f64::INFINITY, |a, &b| a.min(b));
    let max_reward = rewards.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));

    // 2. 打印统计信息
    println!("\n训练结果:");
    println!("  训练轮数: {}", rewards.len());
    println!("  平均奖励: {:.4}", avg_reward);
    println!("  最低奖励: {:.4}", min_reward);
    println!("  最高奖励: {:.4}", max_reward);

    // 3. 保存奖励数据到CSV文件
    let csv_path = rewards_path.with_extension("csv");
    if let Ok(mut file) = File::create(&csv_path) {
        writeln!(file, "Episode,Reward").unwrap();
        for (i, &reward) in rewards.iter().enumerate() {
            writeln!(file, "{},{}", i + 1, reward).unwrap();
        }
        println!("\n奖励数据已保存至: {:?}", csv_path);
    } else {
        eprintln!("无法创建奖励数据文件: {:?}", csv_path);
    }

    // 4. 生成奖励曲线图
    let png_path = rewards_path.with_extension("png");
    if let Err(e) = plot_rewards_chart(rewards, &png_path) {
        eprintln!("无法生成奖励曲线图: {}", e);
    } else {
        println!("奖励曲线图已保存至: {:?}", png_path);
    }
}

/// 使用plotters库绘制奖励曲线图
#[allow(unused_variables)]
fn plot_rewards_chart(
    rewards: &[f64],
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    // 计算统计值
    let min_reward = *rewards
        .iter()
        .min_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(&0.0);
    let max_reward = *rewards
        .iter()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(&1.0);
    let avg_reward: f64 = rewards.iter().sum::<f64>() / rewards.len() as f64;

    // 添加统计信息文本
    let stats_text = format!(
        "统计信息:\n轮次: {}\n平均奖励: {:.4}\n最高奖励: {:.4}\n最低奖励: {:.4}",
        rewards.len(),
        avg_reward,
        max_reward,
        min_reward
    );

    println!("{}", stats_text);

    // 创建绘图区域
    // let root = BitMapBackend::new(output_path, (1200, 800)).into_drawing_area();
    // root.fill(&WHITE)?;
    //
    // let root = root.margin(10, 10, 10, 10);

    // // 创建图表上下文
    // let mut chart = ChartBuilder::on(&root)
    //     .caption("训练奖励曲线", ("sans-serif", 30).into_font())
    //     .x_label_area_size(40)
    //     .y_label_area_size(50)
    //     .margin(5)
    //     .build_cartesian_2d(
    //         0..rewards.len() as u32,
    //         (min_reward - 0.1)..(max_reward + 0.1),
    //     )?;
    //
    // // 配置网格和轴标签
    // chart
    //     .configure_mesh()
    //     .x_desc("训练轮次 (Episode)")
    //     .y_desc("奖励 (Reward)")
    //     .axis_desc_style(("sans-serif", 20))
    //     .x_labels(15)
    //     .y_labels(10)
    //     .draw()?;
    //
    // // 绘制原始奖励曲线 (蓝色)
    // chart.draw_series(LineSeries::new(
    //     rewards.iter().enumerate().map(|(i, r)| (i as u32, *r)),
    //     &BLUE,
    // ))?;
    //
    // // 绘制移动平均线 (红色)
    // let window_size = (rewards.len() / 50).clamp(10, 100);
    // let moving_avg: Vec<f64> = rewards
    //     .windows(window_size)
    //     .map(|w| w.iter().sum::<f64>() / w.len() as f64)
    //     .collect();
    //
    // chart.draw_series(LineSeries::new(
    //     moving_avg
    //         .iter()
    //         .enumerate()
    //         .map(|(i, r)| ((i + window_size - 1) as u32, *r)),
    //     &RED,
    // ))?;
    //
    // // 添加图例
    // chart
    //     .configure_series_labels()
    //     .background_style(WHITE.mix(0.8))
    //     .border_style(BLACK)
    //     .position(SeriesLabelPosition::UpperRight)
    //     .label_font(("sans-serif", 18))
    //     .draw()?;
    //
    // root.draw(&Text::new(
    //     stats_text,
    //     (750, 30),
    //     ("sans-serif", 18).into_font(),
    // ))?;
    //
    // // 添加图例说明
    // root.draw(&Text::new(
    //     "原始奖励",
    //     (750, 60),
    //     ("sans-serif", 18).into_font().color(&BLUE),
    // ))?;
    //
    // root.draw(&Text::new(
    //     format!("{}-轮移动平均", window_size),
    //     (750, 90),
    //     ("sans-serif", 18).into_font().color(&RED),
    // ))?;

    Ok(())
}
